{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Title: #Maximum Points After Collecting Coins From All Nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Difficulty: #Hard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Category Title: #Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tag Slug: #bit-manipulation #tree #depth-first-search #array #dynamic-programming"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Name Translated: #位运算 #树 #深度优先搜索 #数组 #动态规划"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solution Name: maximumPoints"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Title: #收集所有金币可获得的最大积分"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Content:\n",
    "<p>有一棵由 <code>n</code> 个节点组成的无向树，以&nbsp;<code>0</code>&nbsp; 为根节点，节点编号从 <code>0</code> 到 <code>n - 1</code> 。给你一个长度为 <code>n - 1</code> 的二维 <strong>整数</strong> 数组 <code>edges</code> ，其中 <code>edges[i] = [a<sub>i</sub>, b<sub>i</sub>]</code> 表示在树上的节点 <code>a<sub>i</sub></code> 和 <code>b<sub>i</sub></code> 之间存在一条边。另给你一个下标从 <strong>0</strong> 开始、长度为 <code>n</code> 的数组 <code>coins</code> 和一个整数 <code>k</code> ，其中 <code>coins[i]</code> 表示节点 <code>i</code> 处的金币数量。</p>\n",
    "\n",
    "<p>从根节点开始，你必须收集所有金币。要想收集节点上的金币，必须先收集该节点的祖先节点上的金币。</p>\n",
    "\n",
    "<p>节点 <code>i</code> 上的金币可以用下述方法之一进行收集：</p>\n",
    "\n",
    "<ul>\n",
    "\t<li>收集所有金币，得到共计 <code>coins[i] - k</code> 点积分。如果 <code>coins[i] - k</code> 是负数，你将会失去 <code>abs(coins[i] - k)</code> 点积分。</li>\n",
    "\t<li>收集所有金币，得到共计 <code>floor(coins[i] / 2)</code> 点积分。如果采用这种方法，节点 <code>i</code> 子树中所有节点 <code>j</code> 的金币数 <code>coins[j]</code> 将会减少至 <code>floor(coins[j] / 2)</code> 。</li>\n",
    "</ul>\n",
    "\n",
    "<p>返回收集 <strong>所有</strong> 树节点的金币之后可以获得的最大积分。</p>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong class=\"example\">示例 1：</strong></p>\n",
    "<img alt=\"\" src=\"https://assets.leetcode.com/uploads/2023/09/18/ex1-copy.png\" style=\"width: 60px; height: 316px; padding: 10px; background: rgb(255, 255, 255); border-radius: 0.5rem;\" />\n",
    "<pre>\n",
    "<strong>输入：</strong>edges = [[0,1],[1,2],[2,3]], coins = [10,10,3,3], k = 5\n",
    "<strong>输出：</strong>11                        \n",
    "<strong>解释：</strong>\n",
    "使用第一种方法收集节点 0 上的所有金币。总积分 = 10 - 5 = 5 。\n",
    "使用第一种方法收集节点 1 上的所有金币。总积分 = 5 + (10 - 5) = 10 。\n",
    "使用第二种方法收集节点 2 上的所有金币。所以节点 3 上的金币将会变为 floor(3 / 2) = 1 ，总积分 = 10 + floor(3 / 2) = 11 。\n",
    "使用第二种方法收集节点 3 上的所有金币。总积分 =  11 + floor(1 / 2) = 11.\n",
    "可以证明收集所有节点上的金币能获得的最大积分是 11 。 \n",
    "</pre>\n",
    "\n",
    "<p><strong class=\"example\">示例 2：</strong></p>\n",
    "<strong class=\"example\"> <img alt=\"\" src=\"https://assets.leetcode.com/uploads/2023/09/18/ex2.png\" style=\"width: 140px; height: 147px; padding: 10px; background: #fff; border-radius: .5rem;\" /></strong>\n",
    "\n",
    "<pre>\n",
    "<strong>输入：</strong>edges = [[0,1],[0,2]], coins = [8,4,4], k = 0\n",
    "<strong>输出：</strong>16\n",
    "<strong>解释：</strong>\n",
    "使用第一种方法收集所有节点上的金币，因此，总积分 = (8 - 0) + (4 - 0) + (4 - 0) = 16 。\n",
    "</pre>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>提示：</strong></p>\n",
    "\n",
    "<ul>\n",
    "\t<li><code>n == coins.length</code></li>\n",
    "\t<li><code>2 &lt;= n &lt;= 10<sup>5</sup></code></li>\n",
    "\t<li><code><font face=\"monospace\">0 &lt;= coins[i] &lt;= 10<sup>4</sup></font></code></li>\n",
    "\t<li><code>edges.length == n - 1</code></li>\n",
    "\t<li><code><font face=\"monospace\">0 &lt;= edges[i][0], edges[i][1] &lt; n</font></code></li>\n",
    "\t<li><code><font face=\"monospace\">0 &lt;= k &lt;= 10<sup>4</sup></font></code></li>\n",
    "</ul>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Description: [maximum-points-after-collecting-coins-from-all-nodes](https://leetcode.cn/problems/maximum-points-after-collecting-coins-from-all-nodes/description/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solutions: [maximum-points-after-collecting-coins-from-all-nodes](https://leetcode.cn/problems/maximum-points-after-collecting-coins-from-all-nodes/solutions/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_cases = ['[[0,1],[1,2],[2,3]]\\n[10,10,3,3]\\n5', '[[0,1],[0,2]]\\n[8,4,4]\\n0']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        g = defaultdict(list)\n",
    "        n = len(coins)\n",
    "        if k == 0:\n",
    "            return sum(coins)\n",
    "        \n",
    "        for a, b in edges:\n",
    "            g[a].append(b)\n",
    "            g[b].append(a)\n",
    "            \n",
    "        @cache    \n",
    "        def dp(cur, par, p):\n",
    "            if p == 14:\n",
    "                return 0\n",
    "#             ans = 0\n",
    "#             for nei in g[cur]:\n",
    "#                 if nei == par: continue\n",
    "                    \n",
    "#                 ans = max(ans, (coins[cur] >> p) - k + dp(nei, cur, p), coins[cur] >> (p+1) + dp(nei, cur, p+1))\n",
    "                \n",
    "            return max(((coins[cur] >> p) - k) + sum(dp(nei, cur, p) for nei in g[cur] if nei != par), (coins[cur] >> (p+1)) + sum(dp(nei, cur, p+1) for nei in g[cur] if nei != par))\n",
    "        \n",
    "        \n",
    "        return dp(0, -1, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# https://leetcode.cn/u/l00/\n",
    "\n",
    "# https://leetcode.cn/problems/maximum-points-after-collecting-coins-from-all-nodes/solutions/2509808/bo-yang-cong-java-100-shu-wei-zhuang-tai-nib0/\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        n = len(edges) + 1\n",
    "        \n",
    "        deg = [0] * n\n",
    "        deg[0] = 2\n",
    "        nodeMix = [0] * n\n",
    "        for u, v in edges:\n",
    "            deg[u] += 1\n",
    "            deg[v] += 1\n",
    "            nodeMix[u] ^= v\n",
    "            nodeMix[v] ^= u\n",
    "\n",
    "        queue = []\n",
    "        queueLeft = 0\n",
    "        for i, d in enumerate(deg):\n",
    "            if d  == 1: queue.append(i)\n",
    "\n",
    "        # 1 << 14 == 16384\n",
    "        # coinBitlength = 14\n",
    "\n",
    "        coinBitlength = 0\n",
    "        maxCoin = max(coins) << 1\n",
    "        while maxCoin := maxCoin >> 1: coinBitlength += 1\n",
    "        coinBitlength = max(coinBitlength, 2)\n",
    "\n",
    "        dps = [[0] * (coinBitlength + 1) for _ in range(n)]\n",
    "\n",
    "        while queueLeft < len(queue):\n",
    "            curNode = queue[queueLeft]\n",
    "            parent = nodeMix[curNode]\n",
    "            nodeMix[parent] ^= curNode\n",
    "            deg[parent] -= 1\n",
    "            if deg[parent] == 1: queue.append(parent)\n",
    "\n",
    "            res = dps[curNode]\n",
    "            coin = coins[curNode]\n",
    "            for i in range(coinBitlength):\n",
    "                theK = (coin >> i) - k + res[i]\n",
    "                theHalf = (coin >> (i + 1)) + res[i + 1]\n",
    "                res[i] = max(theHalf, theK)\n",
    "                dps[parent][i] += res[i]\n",
    "\n",
    "            queueLeft += 1\n",
    "\n",
    "        coin = coins[0]\n",
    "        return max(coin - k + dps[0][0], (coin >> 1) + dps[0][1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        N = len(coins)\n",
    "\n",
    "        # 每个点的边\n",
    "        de = defaultdict(list)\n",
    "\n",
    "        for a, b in edges:\n",
    "            de[a].append(b)\n",
    "            de[b].append(a)\n",
    "\n",
    "        # 每个节点每一层最大值，计算到节点层数或者节点值等于0为止\n",
    "        lms = [0] * N\n",
    "\n",
    "        def get(lm, i):\n",
    "            if i >= len(lm):\n",
    "                return 0\n",
    "            return lm[i]\n",
    "\n",
    "        # 使用bfs\n",
    "        dp = [(0, 0)]\n",
    "\n",
    "        while dp:\n",
    "            a, c = dp[-1]\n",
    "            if lms[a] == 0:\n",
    "                lms[a] = []\n",
    "\n",
    "                for b in de[a]:\n",
    "                    if lms[b] == 0:\n",
    "                        dp.append((b, c + 1))\n",
    "            else:\n",
    "                dp.pop()\n",
    "                lma = lms[a]\n",
    "                ca = coins[a]\n",
    "                i = 0\n",
    "\n",
    "                while 1:\n",
    "                    # 两个方案都看一下\n",
    "                    ret1 = ca - k\n",
    "                    ret2 = ca // 2\n",
    "\n",
    "                    for b in de[a]:\n",
    "                        lmb = lms[b]\n",
    "                        ret1 += get(lmb, i)\n",
    "                        ret2 += get(lmb, i + 1)\n",
    "\n",
    "                    ret = max(ret1, ret2)\n",
    "                    if ret == 0:\n",
    "                        break\n",
    "\n",
    "                    lma.append(ret)\n",
    "\n",
    "                    if i == c:\n",
    "                        break\n",
    "                    else:\n",
    "                        i += 1\n",
    "                        ca //= 2\n",
    "\n",
    "        return get(lms[0], 0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# -*- coding: utf-8 -*-\n",
    "from typing import List, Tuple, Optional\n",
    "from collections import deque, Counter\n",
    "from queue import PriorityQueue\n",
    "import math\n",
    "from functools import lru_cache\n",
    "from sortedcontainers import SortedDict, SortedSet\n",
    "import random\n",
    "import copy\n",
    "\n",
    "import sys\n",
    "sys.setrecursionlimit(9999999)\n",
    "\n",
    "MOD = 10**9 + 7\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "\n",
    "        n = len(coins)\n",
    "        lk = [[] for i in range(n)]\n",
    "        for a, b in edges:\n",
    "            lk[a].append(b)\n",
    "            lk[b].append(a)\n",
    "\n",
    "\n",
    "\n",
    "        que = deque()\n",
    "        que.append(0)\n",
    "        vis = set()\n",
    "        vis.add(0)\n",
    "        t = []\n",
    "\n",
    "        ch = [[] for i in range(n)]\n",
    "        while len(que):\n",
    "            x = que.popleft()\n",
    "            t.append(x)\n",
    "\n",
    "            for y in lk[x]:\n",
    "                if y not in vis:\n",
    "                    vis.add(y)\n",
    "                    que.append(y)\n",
    "                    ch[x].append(y)\n",
    "\n",
    "\n",
    "        dp = [[0]* 15 for i in range(n)]\n",
    "        for x in t[::-1]:\n",
    "\n",
    "            dp[x][14] = 0\n",
    "            for div in range(14):\n",
    "                cur = coins[x] >> div\n",
    "\n",
    "                ans1 = cur-k\n",
    "                for y in ch[x]:\n",
    "                    ans1 += dp[y][div]\n",
    "\n",
    "\n",
    "                ans2 = cur >> 1\n",
    "                for y in ch[x]:\n",
    "                    ans2 += dp[y][div+1]\n",
    "\n",
    "                dp[x][div] = max(ans1, ans2)\n",
    "\n",
    "        return dp[0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# https://leetcode.cn/u/l00/\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        n = len(edges) + 1\n",
    "        g = [[] for _ in range(n)]\n",
    "        \n",
    "        for edge in edges:\n",
    "            g[edge[0]].append(edge[1])\n",
    "            g[edge[1]].append(edge[0])\n",
    "\n",
    "        # 1 << 14 == 16384\n",
    "        # coinBitlength = 14\n",
    "\n",
    "        coinBitlength = 0\n",
    "        maxCoin = max(coins) << 1\n",
    "        while maxCoin := maxCoin >> 1: coinBitlength += 1\n",
    "        \n",
    "        \n",
    "        def dfs(root: int, pre: int) -> List[int]:\n",
    "            res = [0] * (coinBitlength + 1)\n",
    "            for nex in g[root]:\n",
    "                if nex == pre: continue\n",
    "                nextRes = dfs(nex, root)\n",
    "                for i, num in enumerate(nextRes):\n",
    "                    res[i] += num\n",
    "            coin = coins[root]\n",
    "            for i in range(coinBitlength):\n",
    "                theK = (coin >> i) - k + res[i]\n",
    "                theHalf = (coin >> (i + 1)) + res[i + 1]\n",
    "                res[i] = max(theHalf, theK)\n",
    "                # if res[i] == 0: break\n",
    "            return res\n",
    "        \n",
    "        return dfs(0, -1)[0]\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        n=len(coins);edge,f=[[] for i in range(n)],[[0]*15 for i in range(n)]\n",
    "        for i in edges:edge[i[0]].append(i[1]);edge[i[1]].append(i[0])\n",
    "        def dfs(u,v):\n",
    "            for i in edge[u]:\n",
    "                if(i!=v):dfs(i,u)\n",
    "            g=[0]*15\n",
    "            for i in range(15):\n",
    "                f[u][i]=coins[u]>>i\n",
    "                for j in edge[u]:\n",
    "                    if(j!=v):g[i]+=f[j][i]\n",
    "            for i in range(14):f[u][i]=max(f[u][i]-k+g[i],(f[u][i]>>1)+g[i+1])\n",
    "        dfs(0,0)\n",
    "        return f[0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        N = len(edges) + 1\n",
    "        graph = [[] for _ in range(N)]\n",
    "        for u, v in edges:\n",
    "            graph[u].append(v)\n",
    "            graph[v].append(u)\n",
    "\n",
    "        def dfs(x, p):\n",
    "            bestk = coins[x] - k\n",
    "            c = coins[x]\n",
    "            half = []\n",
    "            while c >= 2:\n",
    "                c //= 2\n",
    "                half.append(c)\n",
    "            for b in graph[x]:\n",
    "                if b != p:\n",
    "                    b, h = dfs(b, x)\n",
    "                    bestk += b\n",
    "                    for i, v in enumerate(h):\n",
    "                        if i >= len(half):\n",
    "                            half.append(0)\n",
    "                        half[i] += v\n",
    "            if not half:\n",
    "                return max(bestk, 0), []\n",
    "            else:\n",
    "                best = max(bestk, half[0])\n",
    "                return best, half[1:]\n",
    "\n",
    "        return dfs(0, None)[0]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        N = len(edges) + 1\n",
    "        graph = [[] for _ in range(N)]\n",
    "        for u, v in edges:\n",
    "            graph[u].append(v)\n",
    "            graph[v].append(u)\n",
    "\n",
    "        def dfs(x, p, d):\n",
    "            bestk = coins[x] - k\n",
    "            c = coins[x]\n",
    "            half = []\n",
    "            while c >= 2 and len(half) <= d:\n",
    "                c //= 2\n",
    "                half.append(c)\n",
    "            for b in graph[x]:\n",
    "                if b != p:\n",
    "                    b, h = dfs(b, x, d + 1)\n",
    "                    bestk += b\n",
    "                    for i, v in enumerate(h):\n",
    "                        if i >= len(half):\n",
    "                            half.append(0)\n",
    "                        half[i] += v\n",
    "            if not half:\n",
    "                return max(bestk, 0), []\n",
    "            else:\n",
    "                best = max(bestk, half[0])\n",
    "                return best, half[1:]\n",
    "\n",
    "        return dfs(0, None, 0)[0]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution(object):\n",
    "    def maximumPoints(self, edges, coins, k):\n",
    "        \"\"\"\n",
    "        :type edges: List[List[int]]\n",
    "        :type coins: List[int]\n",
    "        :type k: int\n",
    "        :rtype: int\n",
    "        \"\"\"\n",
    "\n",
    "        n = len(edges) + 1\n",
    "        g = {index : set() for index in range(n)}\n",
    "        for x, y in edges:\n",
    "            g[x].add(y)\n",
    "            g[y].add(x)\n",
    "\n",
    "        f = [[None] * 14 for _ in range(n)]\n",
    "\n",
    "        def dfs(now, right_shift, fa):\n",
    "            if right_shift == 14:\n",
    "                return 0\n",
    "\n",
    "            if f[now][right_shift] is None:\n",
    "                v1 = (coins[now] >> right_shift) - k\n",
    "                v2 = coins[now] >> (right_shift + 1)\n",
    "                for child in g[now]:\n",
    "                    if child == fa:\n",
    "                        continue\n",
    "                    \n",
    "                    v1 += dfs(child, right_shift, now)\n",
    "                    v2 += dfs(child, right_shift + 1, now)\n",
    "            \n",
    "                f[now][right_shift] = max(v1, v2)\n",
    "            \n",
    "            return f[now][right_shift]\n",
    "        \n",
    "        return dfs(0, 0, None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from collections import deque\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        g = [[] for _ in coins]\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "        cache = [[-1 for _ in range(14)] for i in range(len(coins))]\n",
    "\n",
    "        def dfs(i, j, fa, cache):\n",
    "            if cache[i][j] != -1:\n",
    "                return cache[i][j]\n",
    "            res1 = (coins[i] >> j) - k\n",
    "            res2 = coins[i] >> (j + 1)\n",
    "            for ch in g[i]:\n",
    "                if ch != fa:\n",
    "                    res1 += dfs(ch, j, i, cache)\n",
    "                    if j < 13:\n",
    "                        res2 += dfs(ch, j + 1, i, cache)\n",
    "\n",
    "            cache[i][j] = max(res1, res2)\n",
    "            return cache[i][j]\n",
    "        return dfs(0,0,-1, cache)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        n = len(coins)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for a, b in edges:\n",
    "            g[a].append(b)\n",
    "            g[b].append(a)\n",
    "        # floor(coins[i] / 2)操作就是整除2，整除2就是相当于右移一位\n",
    "        # optimization:1e4最多被整除14次就变成0了，所以这个状态个数最大14\n",
    "        # dp[i][j]表示节点u为根的子树中，所有字节的被整除了j次，的方案的最大值\n",
    "        dp = [[-1] * 14 for _ in range(n)]\n",
    "\n",
    "        def dfs(u, fa, cnt):\n",
    "            if dp[u][cnt] != -1:\n",
    "                return dp[u][cnt]\n",
    "            res1, res2 = (coins[u] >> cnt) - k, (coins[u] >> (cnt + 1))\n",
    "\n",
    "            for j in g[u]:\n",
    "                if j != fa:\n",
    "                    res1 += dfs(j, u, cnt)\n",
    "                    if cnt + 1 < 14:\n",
    "                        res2 += dfs(j, u, cnt + 1)\n",
    "            dp[u][cnt] = max(res1, res2)\n",
    "            return dp[u][cnt]\n",
    "\n",
    "        return dfs(0, -1, 0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        n=len(coins)\n",
    "        tree=[[] for _ in range(n)]\n",
    "        fa=[-1]*n\n",
    "        for u,v in edges:\n",
    "            tree[u].append(v)\n",
    "            tree[v].append(u)\n",
    "        q=[0]\n",
    "        while len(q)>0:\n",
    "            t=q[0]\n",
    "            for child in tree[t]:\n",
    "                if child!=fa[t]:\n",
    "                    fa[child]=t\n",
    "                    q.append(child)\n",
    "            q.pop(0)\n",
    "        # print(fa)\n",
    "        dp=[[-1 for _ in range(20)] for _ in range(n)]\n",
    "        \n",
    "        def helper(co, t):\n",
    "            for _ in range(t):\n",
    "                co = floor(co / 2)\n",
    "            return co\n",
    "        \n",
    "        def getcoins(n, t):\n",
    "            if t>15: return\n",
    "            if dp[n][t]!=-1: return\n",
    "            \n",
    "            co = helper(coins[n], t)\n",
    "            \n",
    "            res1 = co - k\n",
    "            for child in tree[n]:\n",
    "                if child!=fa[n]:\n",
    "                    getcoins(child, t)\n",
    "                    if t<=15:\n",
    "                        res1+=dp[child][t]\n",
    "            \n",
    "            res2 = floor(co / 2)\n",
    "            for child in tree[n]:\n",
    "                if child!=fa[n]:\n",
    "                    getcoins(child, t+1)\n",
    "                    if t+1<=15:\n",
    "                        res2+=dp[child][t+1]\n",
    "                \n",
    "            res=max(res1,res2)\n",
    "            dp[n][t]=res\n",
    "            \n",
    "        \n",
    "        getcoins(0, 0)\n",
    "        return dp[0][0]\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        n = len(coins)\n",
    "        G = [[] for _ in range(n)]\n",
    "        for a,b in edges:\n",
    "            G[a].append(b)\n",
    "            G[b].append(a)\n",
    "        m = max(coins).bit_length()\n",
    "        \n",
    "        # @cache\n",
    "        # def dfs(i, t, p):\n",
    "        #     if t >= m:\n",
    "        #         return 0\n",
    "        #     ans1 = (coins[i] >> t) - k\n",
    "        #     ans2 = coins[i] >> (t+1)\n",
    "        #     for j in G[i]:\n",
    "        #         if j == p: continue\n",
    "        #         ans1 += dfs(j, t, i)\n",
    "        #         ans2 += dfs(j, t+1, i)\n",
    "        #     return max(ans1, ans2)\n",
    "        # return dfs(0, 0, -1)\n",
    "        \n",
    "        f = [[0]*(m+1) for _ in range(n)]\n",
    "        def dfs(i, p):\n",
    "            g = [0] * m\n",
    "            for t in range(m):\n",
    "                f[i][t] = (coins[i] >> t) - k\n",
    "                g[t] = coins[i] >> (t+1)\n",
    "            for j in G[i]:\n",
    "                if j == p: continue\n",
    "                dfs(j, i)\n",
    "                for t in range(m):\n",
    "                    f[i][t] += f[j][t]\n",
    "                    g[t] += f[j][t+1]\n",
    "            for t in range(m):\n",
    "                f[i][t] = max(f[i][t], g[t])\n",
    "        dfs(0, -1)\n",
    "        return f[0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        g = [[] for _ in range(len(coins))]\n",
    "        for f, t in edges:\n",
    "            g[f].append(t)\n",
    "            g[t].append(f)\n",
    "\n",
    "        def dfs(i, fa):\n",
    "            res1 = [(coins[i] >> j) - k for j in range(14)]\n",
    "            res2 = [coins[i] >> (j + 1) for j in range(14)]\n",
    "            for ch in g[i]:\n",
    "                if ch == fa:\n",
    "                    continue\n",
    "                r = dfs(ch, i)\n",
    "                for j in range(14):\n",
    "                    res1[j] += r[j]\n",
    "                    if j + 1 < 14:\n",
    "                        res2[j] += r[j + 1]\n",
    "            for j in range(14):\n",
    "                res1[j] = max(res1[j], res2[j])\n",
    "            return res1\n",
    "        return dfs(0, -1)[0]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        n = len(coins)\n",
    "        to = [[] for _ in range(n)]\n",
    "        for u, v in edges:\n",
    "            to[u].append(v)\n",
    "            to[v].append(u)\n",
    "        f = [[0] * 15 for _ in range(n)]\n",
    "        g = [[0] * 15 for _ in range(n)]\n",
    "        \n",
    "        def dfs(u, fa):\n",
    "            for i in range(15): \n",
    "                f[u][i] = coins[u] - k\n",
    "                g[u][i] = coins[u] // (1 << i) // 2\n",
    "            for v in to[u]:\n",
    "                if v != fa:\n",
    "                    dfs(v, u)\n",
    "                    for i in range(15):\n",
    "                        f[u][i] += max(f[v][i], g[v][i])\n",
    "                        if i < 14:\n",
    "                            g[u][i] += g[v][i + 1]\n",
    "        dfs(0, -1)\n",
    "        return max(f[0][0], g[0][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\"\"\"Created: 2023/10/29\"\"\"\n",
    "from typing import List, Dict\n",
    "from collections import defaultdict\n",
    "from functools import cache\n",
    "\n",
    "\n",
    "class Node:\n",
    "    def __init__(self, v, id_):\n",
    "        self.v = v\n",
    "        self.id = id_\n",
    "        self.children = []\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        n = len(coins)\n",
    "        connected = defaultdict(list)\n",
    "        for e in edges:\n",
    "            connected[e[0]].append(e[1])\n",
    "            connected[e[1]].append(e[0])\n",
    "\n",
    "        def build_tree(node, parent):\n",
    "            for c in connected[node.id]:\n",
    "                if c == parent:\n",
    "                    continue\n",
    "                child = Node(coins[c], c)\n",
    "                node.children.append(child)\n",
    "                build_tree(child, node.id)\n",
    "        root = Node(coins[0], 0)\n",
    "        build_tree(root, -1)\n",
    "\n",
    "        table =[]\n",
    "        for _ in range(n):\n",
    "            table.append({})\n",
    "        def dfs(node, divide):\n",
    "            origin_divide = divide\n",
    "            if divide in table[node.id]:\n",
    "                return table[node.id][divide]\n",
    "            if divide // 2 in table[node.id] and table[node.id][divide // 2] == 0:\n",
    "                table[node.id][divide] = 0\n",
    "                return 0\n",
    "            ans1 = node.v // divide - k\n",
    "            for c in node.children:\n",
    "                ans1 += dfs(c, divide)\n",
    "            divide *= 2\n",
    "            ans2 = node.v // divide\n",
    "            for c in node.children:\n",
    "                ans2 += dfs(c, divide)\n",
    "            ans = max(ans1, ans2)\n",
    "            table[node.id][origin_divide] = ans\n",
    "            return ans\n",
    "        aa = dfs(root, 1)\n",
    "        return aa\n",
    "\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    s = Solution()\n",
    "    print(s.maximumPoints(\n",
    "        [[0, 1], [2, 0], [0, 3], [4, 2]],\n",
    "        [7, 5, 0, 9, 3],\n",
    "    4\n",
    "    ))\n",
    "\n",
    "    # print(s.maximumPoints(\n",
    "    #     [[0, 1], [1, 2], [2, 3]],\n",
    "    #     [10, 10, 3, 3],\n",
    "    #     5\n",
    "    # ))\n",
    "    #\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        n = len(coins)\n",
    "        G = [[] for _ in range(n)]\n",
    "        for a,b in edges:\n",
    "            G[a].append(b)\n",
    "            G[b].append(a)\n",
    "        m = max(coins).bit_length()\n",
    "        \n",
    "        # @cache\n",
    "        # def dfs(i, t, p):\n",
    "        #     if t >= m:\n",
    "        #         return 0\n",
    "        #     ans1 = (coins[i] >> t) - k\n",
    "        #     ans2 = coins[i] >> (t+1)\n",
    "        #     for j in G[i]:\n",
    "        #         if j == p: continue\n",
    "        #         ans1 += dfs(j, t, i)\n",
    "        #         ans2 += dfs(j, t+1, i)\n",
    "        #     return max(ans1, ans2)\n",
    "        # return dfs(0, 0, -1)\n",
    "        \n",
    "        f = [[0]*(m+1) for _ in range(n)]\n",
    "        g = [[0]*m for _ in range(n)]\n",
    "        def dfs(i, p):\n",
    "            for t in range(m):\n",
    "                f[i][t] = (coins[i] >> t) - k\n",
    "                g[i][t] = coins[i] >> (t+1)\n",
    "            for j in G[i]:\n",
    "                if j == p: continue\n",
    "                dfs(j, i)\n",
    "                for t in range(m):\n",
    "                    f[i][t] += f[j][t]\n",
    "                    g[i][t] += f[j][t+1]\n",
    "            for t in range(m):\n",
    "                f[i][t] = max(f[i][t], g[i][t])\n",
    "        dfs(0, -1)\n",
    "        return f[0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        n = len(edges)\n",
    "        g = [[] for _ in range(n + 1)]\n",
    "        for a, b in edges:\n",
    "            g[a].append(b)\n",
    "            g[b].append(a)\n",
    "\n",
    "        @cache\n",
    "        def dfs(fa, now, c):\n",
    "            res1 = (coins[now] >> c) - k\n",
    "            res2 = coins[now] >> (c + 1)\n",
    "            for son in g[now]:\n",
    "                if son == fa:\n",
    "                    continue\n",
    "                res1 += dfs(now, son, c)\n",
    "                if c + 1 < 13:  # coin[i]<=10^4，最多右移13次就为0了\n",
    "                    res2 += dfs(now, son, c + 1)\n",
    "            return max(res1, res2)\n",
    "\n",
    "        return dfs(-1, 0, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:#题解方法，同时dp14种情况\n",
    "        n=len(coins)\n",
    "        nbs=[[] for i in range(n)]\n",
    "        for e0,e1 in edges:\n",
    "            nbs[e0].append(e1)\n",
    "            nbs[e1].append(e0)\n",
    "        def dfs(cur,prev):\n",
    "            res=[[0,0] for i in range(14)]\n",
    "            for i in range(14):\n",
    "                res[i]=[coins[cur]-k,coins[cur]>>1]\n",
    "                coins[cur]>>=1\n",
    "                # if coins[cur]==0:\n",
    "                #     break\n",
    "            # res0=res.copy()\n",
    "            for i in nbs[cur]:\n",
    "                if i!=prev:\n",
    "                    tp=dfs(i,cur)\n",
    "                    for i in range(14):\n",
    "                        res[i][0]+=max(tp[i])\n",
    "                        if i<14-1:\n",
    "                            res[i][1]+=max(tp[i+1])\n",
    "            return res\n",
    "        \n",
    "        res=dfs(0,-1)[0]\n",
    "        # print(f\"{res=}\")\n",
    "        return max(res)\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        g = [[] for _ in coins]\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "\n",
    "        @cache\n",
    "        def dfs(i: int, j: int, fa: int) -> int:\n",
    "            res1 = (coins[i] >> j) - k\n",
    "            res2 = coins[i] >> (j + 1)\n",
    "            flag = res1 < res2\n",
    "            for ch in g[i]:\n",
    "                if ch != fa:\n",
    "                    res1 += dfs(ch, j, i)  # 不右移\n",
    "                    if j < 13 and flag:  # j+1 >= 14 相当于 res2 += 0，无需递归\n",
    "                        res2 += dfs(ch, j + 1, i)  # 右移\n",
    "            return max(res1, res2)\n",
    "        return dfs(0, 0, -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        d = defaultdict(list)\n",
    "        for a, b in edges:\n",
    "            d[a].append(b)\n",
    "            d[b].append(a)\n",
    "        vis = defaultdict(bool)\n",
    "\n",
    "        @lru_cache(None)\n",
    "        def f(i, t):\n",
    "            vis[i] = True\n",
    "            op_1 = (coins[i] >> t) - k\n",
    "            op_2 = coins[i] >> (t + 1)\n",
    "            if op_1 >= op_2:\n",
    "                for j in d[i]:\n",
    "                    if not vis[j]:\n",
    "                        op_1 += f(j, t)\n",
    "            elif t < 13:\n",
    "                for j in d[i]:\n",
    "                    if not vis[j]:\n",
    "                        op_1 += f(j, t)\n",
    "                        op_2 += f(j, t + 1)\n",
    "            vis[i] = False\n",
    "            return max(op_1, op_2)\n",
    "        \n",
    "        return f(0, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        if k == 0:\n",
    "            return sum(coins)\n",
    "        n = len(edges) + 1\n",
    "        g = defaultdict(set)\n",
    "        for x, y in edges:\n",
    "            g[x].add(y)\n",
    "            g[y].add(x)\n",
    "        @cache\n",
    "        def dfs(x, fa, cnt):\n",
    "            res1 = (coins[x] >> cnt) - k \n",
    "            res2 = coins[x] >> (cnt+1) \n",
    "            for y in g[x]:\n",
    "                if y != fa:         \n",
    "                    res1 += dfs(y, x, cnt)    \n",
    "                    if cnt + 1 < 14: # 剪枝，因为1e4右移14位就变成0，没必要递归下去\n",
    "                        res2 += dfs(y, x, cnt+1)\n",
    "            return max(res1, res2)\n",
    "        return dfs(0, -1, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    '''\n",
    "    节点i选1：dfs1(i,j) = sum(dfs(soni,j))+coins[i]>>j - k\n",
    "    节点i选2：dfs2(i,j) = sum(dfs(soni,j+1))+coins[i]>>j\n",
    "    两者中取最大值：dfs(i,j) = max(dfs1(i,j)，dfs1(i,j))\n",
    "\n",
    "    终止条件：\n",
    "    \n",
    "    起始条件：\n",
    "    dfs(0,0)\n",
    "    '''\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        n = len(coins)\n",
    "        p = [[] for _ in range(n)]\n",
    "        for edge in edges:\n",
    "            p[edge[0]].append(edge[1])\n",
    "            p[edge[1]].append(edge[0])\n",
    "        @cache\n",
    "        def dfs(i: int,j: int,f:int)->int:\n",
    "            tmp1 = (coins[i]>>j) - k\n",
    "            tmp2 = coins[i]>>(j+1)\n",
    "            for soni in p[i]:\n",
    "                if(f == soni):\n",
    "                    continue\n",
    "                tmp1 += dfs(soni,j,i)\n",
    "                if j < 13:\n",
    "                    tmp2 += dfs(soni,j+1,i)\n",
    "            return max(tmp1,tmp2)\n",
    "        return dfs(0,0,-1)\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        n = len(coins)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for e in edges:\n",
    "            x, y = e[0], e[1]\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "        @cache\n",
    "        def dfs(i, j, fa):\n",
    "            res1 = (coins[i] >> j) - k\n",
    "            res2 = coins[i] >> (j + 1) \n",
    "            for ch in g[i]:\n",
    "                if ch != fa:\n",
    "                    res1 += dfs(ch, j, i)\n",
    "                    if j < 14:\n",
    "                        res2 += dfs(ch, j + 1, i)\n",
    "            return max(res1, res2)\n",
    "        return dfs(0, 0, -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\r\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\r\n",
    "        g = defaultdict(list)\r\n",
    "        for a, b in edges:\r\n",
    "            g[a].append(b)\r\n",
    "            g[b].append(a)\r\n",
    "        \r\n",
    "        @cache\r\n",
    "        def f(node, fa, acc):\r\n",
    "            if acc ==  14:\r\n",
    "                return 0 \r\n",
    "            tmp = (coins[node] >> acc) - k\r\n",
    "            tmp2 = (coins[node] >> acc) // 2\r\n",
    "            for new in g[node]:\r\n",
    "                if new == fa:\r\n",
    "                    continue \r\n",
    "                tmp += f(new, node, acc)\r\n",
    "                tmp2 += f(new, node, acc + 1)\r\n",
    "            return max(tmp, tmp2) \r\n",
    "        \r\n",
    "        return f(0,-1,0)\r\n",
    "\r\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        n = len(coins)\n",
    "        tree = [[] for _ in range(n)]\n",
    "        for [x,y] in edges:\n",
    "            tree[min(x, y)].append(max(x,y))\n",
    "        print(tree)\n",
    "        @cache\n",
    "        def dfs(i, prev):\n",
    "            if prev >= 2**14:\n",
    "                return 0\n",
    "            res0, res1 = coins[i]//prev-k, coins[i]//(2*prev)\n",
    "            if tree[i] == None:\n",
    "                return 0\n",
    "            for j in tree[i]:\n",
    "                res0 += dfs(j, prev)\n",
    "                res1 += dfs(j, 2*prev)\n",
    "            return max(res0, res1)\n",
    "        \n",
    "        ans = max(dfs(0, 1), dfs(0, 2))\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        if k == 0:\n",
    "            return sum(coins)\n",
    "        import collections\n",
    "        graph = collections.defaultdict(list)\n",
    "        for u, v in edges:\n",
    "            graph[u].append(v)\n",
    "            graph[v].append(u)\n",
    "        root = 0\n",
    "\n",
    "        @functools.lru_cache(None)\n",
    "        def dfs_p(u, p, op):\n",
    "            subs = 0\n",
    "            for v in graph[u]:\n",
    "                if v != p:\n",
    "                    subs += dfs_p(v, u, op)\n",
    "            cur = coins[u] >> op\n",
    "            res = cur - k + subs\n",
    "\n",
    "            subs = 0\n",
    "            for v in graph[u]:\n",
    "                if v != p:\n",
    "                    subs += dfs_p(v, u, min(17, op + 1))\n",
    "            res = max(res, (cur >> 1) + subs)\n",
    "            return res\n",
    "\n",
    "        return dfs_p(root, -1, 0)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        graph=defaultdict(list)\n",
    "        for x,y in edges:\n",
    "            graph[x].append(y)\n",
    "            graph[y].append(x)\n",
    "\n",
    "\n",
    "        n=len(coins)\n",
    "        f=[[[-inf]*2 for i in range(20)] for j in range(n)]\n",
    "        def dp(u,fa):\n",
    "            for j in range(15):\n",
    "                f[u][j][0]=coins[u]-k\n",
    "                if j>0: f[u][j][1]=coins[u]\n",
    "                coins[u]>>=1\n",
    "            for v in graph[u]:\n",
    "                if v!=fa:\n",
    "                    dp(v,u)\n",
    "                    for j in range(15):\n",
    "                        best=max(f[v][j][0],f[v][j+1][1])\n",
    "                        f[u][j][0]+=best\n",
    "                        f[u][j][1]+=best\n",
    "\n",
    "        dp(0,-1)\n",
    "        ans=0\n",
    "        for j in range(15):\n",
    "            ans=max(ans,f[0][j][0],f[0][j][1])\n",
    "        return ans\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        m = defaultdict(set)\n",
    "        for f, t in edges:\n",
    "            m[f].add(t)\n",
    "            m[t].add(f)\n",
    "        \n",
    "        def dfs(p, last):\n",
    "            for np in list(m[p]):\n",
    "                if np == last:\n",
    "                    m[p].remove(np)\n",
    "                    continue\n",
    "                dfs(np, p)\n",
    "        dfs(0, -1)\n",
    "\n",
    "        @cache\n",
    "        def dp(p, mm):\n",
    "            r1 = coins[p] // mm - k\n",
    "            r2 = coins[p] // mm // 2\n",
    "            for np in m[p]:        \n",
    "                if r1 < r2 and mm < 2 ** 13:\n",
    "                    r2 += dp(np, 2 * mm)\n",
    "                r1 += dp(np, mm)\n",
    "            return max(r1, r2)\n",
    "        \n",
    "        res = dp(0, 1)\n",
    "        return res\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        graph=defaultdict(list)\n",
    "        for i,j in edges:\n",
    "            graph[i].append(j)\n",
    "            graph[j].append(i)\n",
    "        n=len(coins)\n",
    "        f=[[[0]*2 for i in range(20)] for j in range(n)]\n",
    "        def dfs(u,fa):\n",
    "            for j in range(15):\n",
    "                f[u][j][0]+=(coins[u]>>j)-k\n",
    "                f[u][j][1]+=coins[u]>>(j+1)\n",
    "\n",
    "            for v in graph[u]:\n",
    "                if v!=fa:\n",
    "                    dfs(v,u)\n",
    "                    for j in range(15):\n",
    "                        f[u][j][1]+=max(f[v][j+1][0],f[v][j+1][1])\n",
    "                        f[u][j][0]+=max(f[v][j][0],f[v][j][1])\n",
    "        dfs(0,-1)\n",
    "        return max(f[0][0][0],f[0][0][1])\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        g = collections.defaultdict(list)\n",
    "        for a, b in edges:\n",
    "            g[a].append(b)\n",
    "            g[b].append(a)\n",
    "\n",
    "        @functools.lru_cache(None)\n",
    "        def dfs(u, p, h):\n",
    "            c = coins[u] >> h\n",
    "            res1, res2 = c - k, c >> 1\n",
    "            for v in g[u]:\n",
    "                if v == p:\n",
    "                    continue\n",
    "                res1 += dfs(v, u, h)\n",
    "                if h < 13:\n",
    "                    res2 += dfs(v, u, h + 1)\n",
    "            return max(res1, res2)\n",
    "\n",
    "        return dfs(0, -1, 0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        n=len(edges)+1\n",
    "        maxv=int(1e4+1)\n",
    "        chd=[[] for _ in range(n)]\n",
    "        for a,b in edges:\n",
    "            chd[a].append(b)\n",
    "            chd[b].append(a)\n",
    "        @lru_cache(None)\n",
    "        def f(nd,tm,par):\n",
    "            if tm==maxv:\n",
    "                return 0\n",
    "            r1=coins[nd]//tm\n",
    "            r2=r1//2\n",
    "            r1-=k\n",
    "            tm2=min(tm*2,maxv)\n",
    "            for ci in chd[nd]:\n",
    "                if ci==par:\n",
    "                    continue\n",
    "                r1+=f(ci,tm,nd)\n",
    "                r2+=f(ci,tm2,nd)\n",
    "            # print(nd,tm,max(r1,r2))\n",
    "            return max(r1,r2)\n",
    "        return f(0,1,-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from functools import lru_cache\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        nodenum = len(coins)\n",
    "        maxnum = 10**7+7\n",
    "        graph = [[] for _ in range(nodenum)]\n",
    "        vis = [0 for _ in range(nodenum)]\n",
    "        value = [0 for _ in range(nodenum)]\n",
    "        dp = [[-maxnum for _ in range(15)] for _ in range(nodenum)]\n",
    "        for edge in edges:\n",
    "            pot0 = edge[0]\n",
    "            pot1 = edge[1]\n",
    "            graph[pot0].append(pot1)\n",
    "            graph[pot1].append(pot0)\n",
    "        @lru_cache(None)\n",
    "        def dfs(now, depth, fa):\n",
    "            res1 = (coins[now] >> depth) - k\n",
    "            res2 = coins[now] >> (depth+1)\n",
    "            for nxt in graph[now]:\n",
    "                if nxt == fa:\n",
    "                    continue\n",
    "                res1 += dfs(nxt, depth, now)\n",
    "                if depth < 13:\n",
    "                    res2 += dfs(nxt, depth+1, now)\n",
    "            return max(res1, res2)\n",
    "        return dfs(0, 0, -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        n = len(coins)\n",
    "        graph = [[] for _ in range(n)]\n",
    "        for fr, to in edges:\n",
    "            graph[fr].append(to)\n",
    "            graph[to].append(fr)\n",
    "        root = [-1] * n\n",
    "        def dfs(x, r):\n",
    "            for y in graph[x]:\n",
    "                if y != r:\n",
    "                    root[y] = x\n",
    "                    dfs(y, x)\n",
    "\n",
    "        @cache\n",
    "        def f(i, j):\n",
    "            if j >= 14:\n",
    "                return 0\n",
    "            res1 = (coins[i] >> j) - k + sum(f(x, j) for x in graph[i] if x != root[i])\n",
    "            res2 = (coins[i] >> (j + 1)) + sum(f(x, j+1) for x in graph[i] if x != root[i])\n",
    "            return max(res1, res2)\n",
    "\n",
    "        dfs(0, -1)\n",
    "        return f(0, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        n = len(edges) + 1\n",
    "        g = [[] for _ in range(n)]\n",
    "        for u,v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        # get tree\n",
    "        tree = [[] for _ in range(n)]\n",
    "        def dfs(u,pa):\n",
    "            for v in g[u]:\n",
    "                if v != pa:\n",
    "                    tree[u].append(v)\n",
    "                    dfs(v,u)\n",
    "        dfs(0,-1)\n",
    "        \n",
    "        from functools import lru_cache\n",
    "        @lru_cache(None)\n",
    "        def f(i,x):\n",
    "            s1 = sum(f(c,x) for c in tree[i]) + (coins[i]>>x)-k\n",
    "            nx = min(x+1,14)\n",
    "            s2 = sum(f(c,nx) for c in tree[i]) + (coins[i]>>nx)\n",
    "            return max(s1,s2)\n",
    "        f.cache_clear()\n",
    "        return f(0,0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        @cache\n",
    "        def search(val, parent, half, is_root):\n",
    "            if half > 2 ** 13:\n",
    "                return 0\n",
    "            tmp = 0\n",
    "            if len(d[val]) == 1 and not is_root:\n",
    "                return max(coins[val] // half - k, coins[val] // (half*2))\n",
    "            # if coins[val] // half >= k:\n",
    "            #     for child in d[val]:\n",
    "            #         if child == parent:\n",
    "            #             continue\n",
    "            #         tmp += search(child, val, half,False)\n",
    "            #     return tmp + coins[val] // half - k\n",
    "            # else:\n",
    "            tmp1 = coins[val] // half - k\n",
    "            tmp2 = coins[val] // (half*2)\n",
    "            for child in d[val]:\n",
    "                if child == parent:\n",
    "                    continue\n",
    "                tmp1 += search(child, val, half, False)\n",
    "                tmp2 += search(child, val, half*2, False)\n",
    "            # tmp2 = coins[val] // (half*2)\n",
    "            # for child in d[val]:\n",
    "            #     if child == parent:\n",
    "            #         continue\n",
    "            #     tmp2 += search(child, val, half*2, False)\n",
    "            return max(tmp1, tmp2)\n",
    "                \n",
    "        min_num = min(coins)\n",
    "        n = len(coins)\n",
    "        # if k <= min_num:\n",
    "        #     return sum(coins) - n * k\n",
    "        # for edge in edges:\n",
    "        #     node\n",
    "        d = defaultdict(list)\n",
    "        for a,b in edges:\n",
    "            d[a].append(b)\n",
    "            d[b].append(a)\n",
    "        ans = search(0, None,1, True)\n",
    "        return ans "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "\n",
    "        n = len(coins)\n",
    "        g = [[] for v in range(n)]\n",
    "        for e in edges:\n",
    "            g[e[0]].append(e[1])\n",
    "            g[e[1]].append(e[0])\n",
    "\n",
    "        @cache\n",
    "        def half(x, times):\n",
    "            return x // int(pow(2, times))\n",
    "\n",
    "        @cache \n",
    "        def dfs(u, fa, i):\n",
    "            c = half(coins[u], i)\n",
    "            s1, s2 = c - k, c // 2\n",
    "            if i >= 20:\n",
    "                for v in g[u]:\n",
    "                    if v == fa: continue\n",
    "                    s1 += dfs(v, u, i)\n",
    "                return s1\n",
    "            for v in g[u]:\n",
    "                if v == fa: continue\n",
    "                s1 += dfs(v, u, i)\n",
    "                s2 += dfs(v, u, i + 1)\n",
    "            return max(s1, s2)\n",
    "        \n",
    "        return dfs(0, -1, 0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# from sortedcontainers import SortedList\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        n = len(coins)\n",
    "\n",
    "        g = [[] for _ in range(n)]\n",
    "\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "\n",
    "        @cache\n",
    "        def dfs(x: int, fa: int, state: int):\n",
    "            if (1 << state) > 10**4:\n",
    "                return 0\n",
    "            op1 = (\n",
    "                coins[x] // pow(2, state)\n",
    "                - k\n",
    "                + sum(dfs(y, x, state) for y in g[x] if y != fa)\n",
    "            )\n",
    "\n",
    "            op2 = coins[x] // pow(2, state + 1) + sum(\n",
    "                dfs(y, x, state + 1) for y in g[x] if y != fa\n",
    "            )\n",
    "\n",
    "            return max(op1, op2)\n",
    "\n",
    "        return dfs(0, -1, 0)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from string import *\n",
    "from re import *\n",
    "from datetime import *\n",
    "from collections import *\n",
    "from heapq import *\n",
    "from bisect import *\n",
    "from copy import *\n",
    "from math import *\n",
    "from random import *\n",
    "from statistics import *\n",
    "from itertools import *\n",
    "from functools import *\n",
    "from operator import *\n",
    "from io import *\n",
    "from sys import *\n",
    "from json import *\n",
    "from builtins import *\n",
    "from typing import *\n",
    "\n",
    "\"\"\"\n",
    "There exists an undirected tree rooted at node 0 with n nodes labeled from 0 to n - 1. You are given\n",
    "a 2D integer array edges of length n - 1, where edges[i] = [ai, bi] indicates that there is an edge\n",
    "between nodes ai and bi in the tree. You are also given a 0-indexed array coins of size n where\n",
    "coins[i] indicates the number of coins in the vertex i, and an integer k.\n",
    "\n",
    "Starting from the root, you have to collect all the coins such that the coins at a node can only be\n",
    "collected if the coins of its ancestors have been already collected.\n",
    "\n",
    "Coins at nodei can be collected in one of the following ways:\n",
    "\n",
    " * Collect all the coins, but you will get coins[i] - k points. If coins[i] - k is negative then you\n",
    "   will lose abs(coins[i] - k) points.\n",
    " * Collect all the coins, but you will get floor(coins[i] / 2) points. If this way is used, then for\n",
    "   all the nodej present in the subtree of nodei, coins[j] will get reduced to floor(coins[j] / 2).\n",
    "\n",
    "Return the maximum points you can get after collecting the coins from all the tree nodes.\n",
    "\n",
    " \n",
    "\n",
    "Example 1:\n",
    "\n",
    "[https://assets.leetcode.com/uploads/2023/09/18/ex1-copy.png]\n",
    "\n",
    "\n",
    "Input: edges = [[0,1],[1,2],[2,3]], coins = [10,10,3,3], k = 5\n",
    "Output: 11                        \n",
    "Explanation: \n",
    "Collect all the coins from node 0 using the first way. Total points = 10 - 5 = 5.\n",
    "Collect all the coins from node 1 using the first way. Total points = 5 + (10 - 5) = 10.\n",
    "Collect all the coins from node 2 using the second way so coins left at node 3 will be floor(3 / 2) = 1. Total points = 10 + floor(3 / 2) = 11.\n",
    "Collect all the coins from node 3 using the second way. Total points = 11 + floor(1 / 2) = 11.\n",
    "It can be shown that the maximum points we can get after collecting coins from all the nodes is 11. \n",
    "\n",
    "\n",
    "Example 2:\n",
    "\n",
    "[https://assets.leetcode.com/uploads/2023/09/18/ex2.png]\n",
    "\n",
    "\n",
    "Input: edges = [[0,1],[0,2]], coins = [8,4,4], k = 0\n",
    "Output: 16\n",
    "Explanation: \n",
    "Coins will be collected from all the nodes using the first way. Therefore, total points = (8 - 0) + (4 - 0) + (4 - 0) = 16.\n",
    "\n",
    "\n",
    " \n",
    "\n",
    "Constraints:\n",
    "\n",
    " * n == coins.length\n",
    " * 2 <= n <= 105\n",
    " * 0 <= coins[i] <= 104\n",
    " * edges.length == n - 1\n",
    " * 0 <= edges[i][0], edges[i][1] < n\n",
    " * 0 <= k <= 104\n",
    "\"\"\"\n",
    "\n",
    "MOD = 1000000007 # 998244353\n",
    "\n",
    "M = 14\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        g = defaultdict(list)\n",
    "        for u, v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "\n",
    "        @cache\n",
    "        def dp(u, t, p=-1):\n",
    "            return 0 if t >= M else max((coins[u] >> t) - k +\n",
    "                                          sum(dp(v, t, u) for v in g[u] if v != p),\n",
    "                                        (coins[u] >> t) // 2 +\n",
    "                                          sum(dp(v, t + 1, u) for v in g[u] if v != p))\n",
    "                \n",
    "            \n",
    "            return max((coins[u] >> t) - k + sum(dp(v, t, u) for v in g[u] if v != p),\n",
    "                       (coins[u] >> t) // 2 + sum(dp(v, t + 1, u) for v in g[u] if v != p))\n",
    "            \n",
    "            res1 = (coins[u] >> t) - k\n",
    "            for v in g[u]:\n",
    "                if v != p:\n",
    "                    res1 += dp(v, t, u)\n",
    "            \n",
    "            res2 = (coins[u] >> t) // 2\n",
    "            for v in g[u]:\n",
    "                if v != p:\n",
    "                    res2 += dp(v, t + 1, u)\n",
    "\n",
    "            return max(res1, res2)\n",
    "            \n",
    "        return dp(0, 0)\n",
    "            \n",
    "testcases = [\n",
    "    [[[0,1],[1,2],[2,3]], [10,10,3,3], 5],\n",
    "    # [[[0,1],[0,2]], [8,4,4], 0],\n",
    "    # [],\n",
    "]\n",
    "\n",
    "s = Solution()\n",
    "func_name = dir(s)[-1]\n",
    "func = getattr(s, func_name)\n",
    "\n",
    "for args in testcases:\n",
    "    print(func(*args))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        dic = {}\n",
    "        n = len(coins)\n",
    "        for e in edges:\n",
    "            s = min(e)\n",
    "            t = max(e)\n",
    "            if s in dic:\n",
    "                dic[s].append(t)\n",
    "            else:\n",
    "                dic[s] = [t]\n",
    "        @cache\n",
    "        def dfs(v,cnt):    \n",
    "            if cnt>=14: \n",
    "                return 0, 0\n",
    "            cv = coins[v]>>cnt\n",
    "            tot1,tot2 = cv-k,cv>>1\n",
    "            if v not in dic:\n",
    "                return tot1,tot2\n",
    "            for u in dic[v]:  \n",
    "                tot1 += max(dfs(u,cnt))\n",
    "                tot2 += max(dfs(u,cnt+1))\n",
    "            return tot1,tot2\n",
    "        return max(dfs(0,0))\n",
    "        \n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        n = len(coins)\n",
    "        g = defaultdict(list)\n",
    "        for u, v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        @cache\n",
    "        def dfs(i: int, x: int, pre):\n",
    "            if x == 17:\n",
    "                return 0\n",
    "            coin = coins[i] >> x\n",
    "            if len(g[i]) == 1 and g[i][0] == pre:\n",
    "                return max(coin - k, coin >> 1)\n",
    "\n",
    "            case1 = coin - k + sum(dfs(nex, x, i) for nex in g[i] if nex != pre)\n",
    "            case2 = (coin >> 1) + sum(dfs(nex, x + 1, i) for nex in g[i] if nex != pre)\n",
    "            return max(case1, case2)\n",
    "\n",
    "        return dfs(0, 0, -1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        neighbors = defaultdict(list)\n",
    "        for a, b in edges:\n",
    "            neighbors[a].append(b)\n",
    "            neighbors[b].append(a)\n",
    "        \n",
    "        visited = set()\n",
    "        \n",
    "        @lru_cache(len(coins) * 14)\n",
    "        def collect_coins(n, halfed_time=0):\n",
    "            if halfed_time >= 14:\n",
    "                return 0\n",
    "            visited.add(n)\n",
    "            children = [neighbor for neighbor in neighbors[n] if neighbor not in visited]\n",
    "            # collect all\n",
    "            collect_all = (coins[n] >> halfed_time) - k + sum(collect_coins(c, halfed_time) for c in children)\n",
    "            \n",
    "            # collect half\n",
    "            collect_half = (coins[n] >> (halfed_time + 1)) + sum(collect_coins(c, halfed_time + 1) for c in children)\n",
    "            visited.remove(n)\n",
    "            return max(0, collect_all, collect_half)\n",
    "    \n",
    "        return collect_coins(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution(object):\n",
    "        def maximumPoints(self, edges, coins, k):\n",
    "            \"\"\"\n",
    "            :type edges: List[List[int]]\n",
    "            :type coins: List[int]\n",
    "            :type k: int\n",
    "            :rtype: int\n",
    "            \"\"\"\n",
    "            # from functools import lru_cache\n",
    "            d = defaultdict(list)\n",
    "            for e in edges:\n",
    "                d[e[0]].append(e[1])\n",
    "                d[e[1]].append(e[0])\n",
    "\n",
    "            @cache\n",
    "            def dfs(i, i_last, x):\n",
    "                # if d[i] == [i_last]:\n",
    "                #     return [(coins[i] >> x) - k, coins[i] >> x // 2]\n",
    "                ans1 = (coins[i] >> x) - k\n",
    "\n",
    "                new_x = min(x + 1,14)\n",
    "                for n in d[i]:\n",
    "                    if n != i_last:\n",
    "                        ans1 += max(dfs(n, i, x))\n",
    "                ans2 = coins[i] >> new_x\n",
    "                for n in d[i]:\n",
    "                    if n != i_last:\n",
    "                        ans2 += max(dfs(n, i, new_x))\n",
    "                return [ans1, ans2]\n",
    "\n",
    "            root = 0\n",
    "            # noroot = False\n",
    "            # for n in d.keys():\n",
    "            #     # print(n)\n",
    "            #     for e in d.keys():\n",
    "            #         if n in d[e]:\n",
    "            #             # print(n,e)\n",
    "            #             noroot = True\n",
    "            #             break\n",
    "            #     if not noroot:\n",
    "            #         root = n\n",
    "            #     noroot = False\n",
    "            # print(root)\n",
    "            ans = max(dfs(root, -1, 0))\n",
    "            dfs.cache_clear()\n",
    "            return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution(object):\n",
    "        def maximumPoints(self, edges, coins, k):\n",
    "            \"\"\"\n",
    "            :type edges: List[List[int]]\n",
    "            :type coins: List[int]\n",
    "            :type k: int\n",
    "            :rtype: int\n",
    "            \"\"\"\n",
    "            # from functools import lru_cache\n",
    "            d = defaultdict(list)\n",
    "            for e in edges:\n",
    "                d[e[0]].append(e[1])\n",
    "                d[e[1]].append(e[0])\n",
    "\n",
    "            @cache\n",
    "            def dfs(i, i_last, x):\n",
    "                # if d[i] == [i_last]:\n",
    "                #     return [(coins[i] >> x) - k, coins[i] >> x // 2]\n",
    "                ans1 = (coins[i] >> x) - k\n",
    "\n",
    "                new_x = min(x + 1,14)\n",
    "                for n in d[i]:\n",
    "                    if n != i_last:\n",
    "                        ans1 += max(dfs(n, i, x))\n",
    "                ans2 = coins[i] >> new_x\n",
    "                for n in d[i]:\n",
    "                    if n != i_last:\n",
    "                        ans2 += max(dfs(n, i, new_x))\n",
    "                return [ans1, ans2]\n",
    "\n",
    "            root = 0\n",
    "            ans = max(dfs(root, -1, 0))\n",
    "            # dfs.cache_clear()\n",
    "            return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumPoints(self, edges: List[List[int]], coins: List[int], k: int) -> int:\n",
    "        n = len(coins)\n",
    "        e = [[] for _ in range(n)]\n",
    "        for u, v in edges:\n",
    "            e[u].append(v)\n",
    "            e[v].append(u)\n",
    "\n",
    "        mx = max(coins)\n",
    "        m = mx.bit_length()\n",
    "        pow2 = [1] * (m + 1)\n",
    "        for i in range(1, m + 1):\n",
    "            pow2[i] = pow2[i - 1] * 2\n",
    "        \n",
    "        @cache\n",
    "        def dfs(u: int, p: int, c: int) -> list[int]:\n",
    "            if c > m:\n",
    "                return [-k, 0]\n",
    "            cur = coins[u] // pow2[c]\n",
    "            ans = [cur - k, cur // 2]\n",
    "            for v in e[u]:\n",
    "                if v == p:\n",
    "                    continue\n",
    "                ans[0] += max(dfs(v, u, c))\n",
    "                ans[1] += max(dfs(v, u, c + 1))\n",
    "            return ans\n",
    "        \n",
    "        return max(dfs(0, -1, 0))"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
